#include "se2.h"

using namespace std;

typedef Eigen::Matrix<double, 6, 1> Vector6d;
typedef Eigen::Matrix<double, 9, 1> Vector9d;
typedef Eigen::Matrix<double, 12, 1> Vector12d;

Eigen::Vector3d odometry0(0, 0, 0);
Eigen::Vector3d odometry1(0.0597, -2.06042e-06, -5.9321e-05);
Eigen::Vector3d odometry2(0.1194, -8.35443e-06, -7.98171e-05);

Eigen::Vector3d pose0(0.00253282, 0.000136343, -0.000189673);
Eigen::Vector3d pose1(0.1193, -1.43682e-06 , -3.87479e-05);

Eigen::Matrix3d transform_cov_to_matrix3d(const double cov[9]) {
    Eigen::Matrix3d matrix;
    for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < 3; ++j) {
            matrix(i, j) = cov[i * 3 + j];
        }
    }
    return matrix;    
}

Eigen::Matrix3d set_observation_cov(const double &cov_x, const double &cov_y, const double &cov_t) {
    Eigen::Matrix3d matrix;
    matrix(0, 0) = cov_x;
    matrix(1, 1) = cov_y;
    matrix(2, 2) = cov_t;
    return matrix;    
}

Eigen::Vector3d update_pose(const Eigen::Vector3d &predict_pose, const Eigen::Vector3d &delta_pose) {
    Se2 predicte_pose_se2 = Se2(predict_pose[0], predict_pose[1], predict_pose[2]);
    Se2 delta_pose_se2 = Se2(delta_pose[0], delta_pose[1], delta_pose[2]);
    Se2 update_pose_se2 = predicte_pose_se2 + delta_pose_se2;
    Eigen::Vector3d updated_pose(update_pose_se2.x, update_pose_se2.y, update_pose_se2.theta);

    return updated_pose;
}

void clear_preintegration_measure(PreSE2 &pre_se2) {
    for(int i = 0; i < 3; i++)
        pre_se2.meas[i] = 0.0;
    for(int i = 0; i < 9; i++)
        pre_se2.cov[i] = 0.0;
}
void odometry_preintegration(PreSE2 &pre_se2, const Se2 &last_odometry, const Se2 &current_odometry) {
    // preintegration
    Eigen::Map<Eigen::Vector3d> meas(pre_se2.meas);
    Se2 odok = current_odometry - last_odometry;
    Eigen::Vector2d odork(odok.x, odok.y);
    Eigen::Matrix2d Phi_ik = Eigen::Rotation2Dd(meas[2]).toRotationMatrix();
    meas.head<2>() += Phi_ik * odork;
    meas[2] += odok.theta;

    // std::cout << " Phi_ik =======: " << Phi_ik << std::endl;
    // std::cout << " odork =======: " << odork << std::endl;
    // std::cout << " odok.theta =======: " << odok.theta << std::endl;
    // std::cout << " meas =======: " << meas.transpose() << std::endl;
    // std::cout << " pre_se2.meas ================= : " << pre_se2.meas[0] << " , " << pre_se2.meas[1] << " , " << pre_se2.meas[2] << std::endl;

    Eigen::Matrix3d Ak = Eigen::Matrix3d::Identity();
    Eigen::Matrix3d Bk = Eigen::Matrix3d::Identity();
    Ak.block<2,1>(0,2) = Phi_ik * Eigen::Vector2d(-odork[1], odork[0]);
    Bk.block<2,2>(0,0) = Phi_ik;
    Eigen::Map<Eigen::Matrix3d, Eigen::RowMajor> Sigmak(pre_se2.cov);
    Eigen::Matrix3d Sigma_vk = Eigen::Matrix3d::Identity();
    // adjust it
    float sigma_x = 1.0e-2; 
    float sigma_y = 1.0e-2;
    float sigma_t = 1.0e-2; 

    Sigma_vk(0,0) = sigma_x * sigma_x;
    Sigma_vk(1,1) = sigma_y * sigma_y;
    Sigma_vk(2,2) = sigma_t * sigma_t;
    Eigen::Matrix3d Sigma_k_1 = Ak * Sigmak * Ak.transpose() + Bk * Sigma_vk * Bk.transpose();
    Sigmak = Sigma_k_1;

    // std::cout << " Sigmak =======: " << Sigmak << std::endl;
}

void predict(const PreSE2 &pre_se2, const Se2 &last_odometry, Eigen::Vector3d &current_odometry) {
    PreSE2 pre_se2_tmp = pre_se2;
    Eigen::Map<Eigen::Vector3d> meas(pre_se2_tmp.meas);
    // Eigen::Vector3d last_odom(last_odometry.x, last_odometry.y, last_odometry.theta);

    // current_odometry = last_odom + meas;
    Se2 delta_pose_se2 = Se2(meas[0], meas[1], meas[2]);
    Se2 update_pose_se2 = last_odometry + delta_pose_se2;
    current_odometry = Eigen::Vector3d(update_pose_se2.x, update_pose_se2.y, update_pose_se2.theta);

    // cout << "meas: " << meas << endl;
    // cout << "current_odometry: " << current_odometry << endl;
}

void preintegration_residual_jacobian(const PreSE2 &pre_se2, const Eigen::Vector3d &last_predict_odometry, const Eigen::Vector3d &curr_predict_odometry, Eigen::Vector3d &residual, Eigen::Matrix3d &jacobianOplusXi, Eigen::Matrix3d &jacobianOplusXj)
{
    PreSE2 pre_se2_tmp = pre_se2;
    Eigen::Map<Eigen::Vector3d> measurement(pre_se2_tmp.meas);

    Eigen::Rotation2D<double> last_rot(last_predict_odometry[2]);
    Eigen::Rotation2D<double> curr_rot(curr_predict_odometry[2]);

    Eigen::Matrix2d Ri = last_rot.toRotationMatrix();
    Eigen::Vector2d ri = last_predict_odometry.head<2>();
    double ai = last_rot.angle();
    double aj = curr_rot.angle();
    Eigen::Vector2d rj = curr_predict_odometry.head<2>();
 
    // 公式（24）
    residual.head<2>() = Ri.transpose() * (rj-ri) - measurement.head<2>();
    residual[2] = aj - ai - measurement[2];

    // compute jacobian matrix
    Eigen::Vector2d rij = rj-ri;
    Eigen::Vector2d rij_x(-rij[1], rij[0]);
    // 公式（25）
    jacobianOplusXi.block<2,2>(0,0) = -Ri.transpose();
    jacobianOplusXi.block<2,1>(0,2) = -Ri.transpose() * rij_x;
    jacobianOplusXi.block<1,2>(2,0).setZero();
    jacobianOplusXi(2,2) = -1;
 
    jacobianOplusXj.setIdentity();
    jacobianOplusXj.block<2,2>(0,0) = Ri.transpose();
}

double Chi2(const Eigen::Vector3d& residual, const Eigen::Matrix3d& information_matrix) {
    return residual.transpose() * information_matrix * residual;
}

void pose_observation_residual_jacobian(const Eigen::Vector3d &obs_odometry, const Eigen::Vector3d &predict_odometry, Eigen::Vector3d &residual, Eigen::Matrix3d &jacobianOplusX) {
    residual = obs_odometry - predict_odometry;
    jacobianOplusX = -Eigen::Matrix3d::Identity();
}

std::vector<Eigen::Vector3d> optimizer(const std::vector<Eigen::Vector3d> &odometry, const std::vector<Eigen::Vector3d> &poses, const std::vector<PreSE2> &pre_se2_vector) {
    int iterations = 100;
    double cost = 0, last_cost = 0;
    int num_poses = poses.size();
    // cout << "num_poses: " << num_poses << endl;

    std::vector<Eigen::Vector3d> result;
    Eigen::Vector3d last_predict_odometry;
    Eigen::Vector3d curr_predict_odometry;

    Eigen::Matrix3d cov_last_pose = set_observation_cov(100.0, 100.0, 100.0);
    Eigen::Matrix3d cov_curr_pose = set_observation_cov(0.05, 0.05, 0.05);
    Eigen::Matrix<double, 9, 9> information_matrix = Eigen::Matrix<double, 9, 9>::Identity();
    information_matrix.block<3,3>(3,3) = cov_last_pose.inverse();
    information_matrix.block<3,3>(6,6) = cov_curr_pose.inverse();

    for (int iter = 0; iter < iterations; iter++) {
        cost = 0;
        Eigen::Matrix<double, 6, 6> H = Eigen::Matrix<double, 6, 6>::Zero();
        Vector6d b = Vector6d::Zero();
        // compute cost
        for (int i = 1; i < num_poses; i++) {
            // get preintegration
            PreSE2 pre_se2 = pre_se2_vector[i];
            // get predict odometry
            if(iter == 0) {
                last_predict_odometry = odometry[i-1];
                curr_predict_odometry = odometry[i];      

                information_matrix.block<3,3>(0,0) = transform_cov_to_matrix3d(pre_se2.cov).inverse();
            }

            // get observation pose
            Eigen::Vector3d last_pose = poses[i-1];
            Eigen::Vector3d curr_pose = poses[i];
            // get odometry preintegration residual and jacobian
            Eigen::Vector3d preint_residual;
            Eigen::Matrix3d jacobianOplusXi;
            Eigen::Matrix3d jacobianOplusXj;
            preintegration_residual_jacobian(pre_se2, last_predict_odometry, curr_predict_odometry, preint_residual, jacobianOplusXi, jacobianOplusXj);
            cost += Chi2(preint_residual, transform_cov_to_matrix3d(pre_se2.cov));

            // get obs residual and jacobian
            Eigen::Vector3d last_residual;
            Eigen::Vector3d curr_residual;
            Eigen::Matrix3d last_jacobianOplusX;
            Eigen::Matrix3d curr_jacobianOplusX;
            pose_observation_residual_jacobian(last_pose, last_predict_odometry, last_residual, last_jacobianOplusX);
            pose_observation_residual_jacobian(curr_pose, curr_predict_odometry, curr_residual, curr_jacobianOplusX);

            cost += Chi2(last_residual, cov_last_pose.inverse());
            cost += Chi2(curr_residual, cov_curr_pose.inverse());

            Vector9d residual;
            residual.block<3,1>(0,0) = preint_residual;
            residual.block<3,1>(3,0) = last_residual;
            residual.block<3,1>(6,0) = curr_residual;

            Eigen::Matrix<double,9,6> jacobian;
            jacobian.block<3,3>(0,0) = jacobianOplusXi;
            jacobian.block<3,3>(0,3) = jacobianOplusXj;
            jacobian.block<3,3>(3,0) = last_jacobianOplusX;
            jacobian.block<3,3>(3,3) = Eigen::Matrix3d::Zero();
            jacobian.block<3,3>(6,0) = Eigen::Matrix3d::Zero();
            jacobian.block<3,3>(6,3) = curr_jacobianOplusX;
            H += jacobian.transpose() * information_matrix * jacobian;
            b += -jacobian.transpose() * information_matrix * residual;
        }

        // solve dx
        Vector6d dx;
        dx = H.ldlt().solve(b);

        if (isnan(dx[0])) {
            cout << "result is nan!" << endl;
            break;
        }

        if (iter > 0 && cost >= last_cost) {
            // 误差增长了，说明近似的不够好
            cout << "cost: " << cost << ", last cost: " << last_cost << endl;
            break;
        }

        last_predict_odometry = update_pose(last_predict_odometry, dx.block<3, 1>(0,0));
        curr_predict_odometry = update_pose(curr_predict_odometry, dx.block<3, 1>(3,0));

        last_cost = cost;

        cout << "iteration == : " << iter << " cost == : " << cout.precision(12) << cost << endl;
        cout << "dx ================== : " << dx.transpose() << endl;
        cout << "last_predict_odometry : " << last_predict_odometry.transpose()  << endl;
        cout << "curr_predict_odometry : " << curr_predict_odometry.transpose() << endl;
    }
    result.push_back(last_predict_odometry);
    result.push_back(curr_predict_odometry);
    return result;
}

int main() {  
    std::vector<Eigen::Vector3d> odometry, predict_odometry;
    odometry.push_back(odometry0);
    odometry.push_back(odometry1);    
    odometry.push_back(odometry2);

    predict_odometry.push_back(odometry0);

    std::vector<Eigen::Vector3d> poses;
    poses.push_back(pose0);
    poses.push_back(pose1);

    PreSE2 pre_se2;
    std::vector<PreSE2> pre_se2_vector;
    clear_preintegration_measure(pre_se2);
	pre_se2_vector.push_back(pre_se2);

    Se2 last_odometry = Se2(odometry0[0], odometry0[1], odometry0[2]);
    Eigen::Vector3d current_odom;
    for(int i = 1; i < odometry.size(); i++) {
    	Eigen::Vector3d odom = odometry[i];
    	Se2 current_odometry = Se2(odom[0], odom[1], odom[2]);
        
    	odometry_preintegration(pre_se2, last_odometry, current_odometry);

        // cout << "odom : " << odom << endl;
    	last_odometry = current_odometry;
    }

    predict(pre_se2, Se2(odometry0[0], odometry0[1], odometry0[2]), current_odom);

    predict_odometry.push_back(current_odom);
    pre_se2_vector.push_back(pre_se2);
    // cout << "current_odom : " << current_odom << endl;
    
    std::vector<Eigen::Vector3d> result = optimizer(predict_odometry, poses, pre_se2_vector);

    for(auto res : result) {
        std::cout << " result : " << res.transpose() << std::endl;
    }
}
